package com.icss.train_sys_back.handler;

import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.stereotype.Component;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    // 存储所有活跃的WebSocket会话
    private static final ConcurrentHashMap<String, WebSocketSession> sessions = new ConcurrentHashMap<>();
    private static final CopyOnWriteArraySet<WebSocketSession> allSessions = new CopyOnWriteArraySet<>();
    
    // JSON处理器
    private final ObjectMapper objectMapper = new ObjectMapper();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 获取用户标识（可以从session属性或参数中获取）
        String userId = getUserIdFromSession(session);

        sessions.put(userId, session);
        allSessions.add(session);

        System.out.println("用户连接: " + userId + ", 当前在线: " + sessions.size());
        
        // 广播用户列表更新
        broadcastUserList();
        
        // 发送连接成功消息给当前用户
        Map<String, Object> welcomeMessage = new HashMap<>();
        welcomeMessage.put("type", "SYSTEM_MESSAGE");
        welcomeMessage.put("content", "连接成功，欢迎加入聊天！");
        sendJsonMessage(session, welcomeMessage);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 处理客户端消息
        String payload = message.getPayload();
        String userId = getUserIdFromSession(session);
        
        System.out.println("收到来自" + userId + "的消息: " + payload);
        
        try {
            // 尝试解析JSON消息
            Map<String, Object> msgData = objectMapper.readValue(payload, Map.class);
            String type = (String) msgData.get("type");
            
            // 根据消息类型处理
            switch (type) {
                case "DIRECT_MESSAGE":
                    handleDirectMessage(userId, msgData);
                    break;
                case "BROADCAST_MESSAGE":
                    handleBroadcastMessage(userId, msgData);
                    break;
                case "HEARTBEAT":
                    handleHeartbeat(session);
                    break;
                default:
                    // 未知消息类型，作为普通文本处理
                    System.out.println("未知消息类型: " + type);
                    break;
            }
        } catch (Exception e) {
            // 非JSON格式消息，作为普通文本处理
            System.out.println("收到非JSON消息: " + payload);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        String userId = getUserIdFromSession(session);
        sessions.remove(userId);
        allSessions.remove(session);
        System.out.println("用户断开: " + userId + ", 剩余在线: " + sessions.size());
        
        // 广播用户列表更新
        broadcastUserList();
    }

    // 向特定用户发送消息
    public void sendToUser(String userId, String message) {
        WebSocketSession session = sessions.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage("用户"+userId + "向您发送消息："+message));
            } catch (IOException e) {
                System.err.println("向用户 " + userId + " 发送消息失败: " + e.getMessage());
                // 清理无效会话
                sessions.remove(userId);
                broadcastUserList();
            }
        }
    }

    // 向所有用户广播消息
    public void broadcast(String message) {
        for (WebSocketSession session : allSessions) {
            try {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage("广播消息:"+message));
                }
            } catch (IOException e) {
                System.err.println("广播消息失败: " + e.getMessage());
                // 清理无效会话
                allSessions.remove(session);
            }
        }
    }
    
    // 广播用户列表
    private void broadcastUserList() {
        try {
            List<String> onlineUsers = getOnlineUsers();
            Map<String, Object> userListMessage = new HashMap<>();
            userListMessage.put("type", "USER_LIST_UPDATE");
            userListMessage.put("users", onlineUsers);
            
            String jsonMessage = objectMapper.writeValueAsString(userListMessage);
            
            for (WebSocketSession session : allSessions) {
                if (session.isOpen()) {
                    session.sendMessage(new TextMessage(jsonMessage));
                }
            }
        } catch (Exception e) {
            System.err.println("广播用户列表失败: " + e.getMessage());
        }
    }
    
    // 处理私信消息
    private void handleDirectMessage(String fromUserId, Map<String, Object> msgData) {
        String toUserId = (String) msgData.get("to");
        String content = (String) msgData.get("content");
        
        if (toUserId != null && content != null) {
            WebSocketSession targetSession = sessions.get(toUserId);
            if (targetSession != null && targetSession.isOpen()) {
                try {
                    Map<String, Object> message = new HashMap<>();
                    message.put("type", "USER_MESSAGE");
                    message.put("from", fromUserId);
                    message.put("content", content);
                    
                    sendJsonMessage(targetSession, message);
                } catch (Exception e) {
                    System.err.println("发送私信失败: " + e.getMessage());
                    // 清理无效会话
                    sessions.remove(toUserId);
                    broadcastUserList();
                }
            }
        }
    }
    
    // 处理广播消息
    private void handleBroadcastMessage(String fromUserId, Map<String, Object> msgData) {
        String content = (String) msgData.get("content");
        
        if (content != null) {
            try {
                Map<String, Object> message = new HashMap<>();
                message.put("type", "USER_MESSAGE");
                message.put("from", fromUserId + "(广播)");
                message.put("content", content);
                
                String jsonMessage = objectMapper.writeValueAsString(message);
                
                for (WebSocketSession session : allSessions) {
                    if (session.isOpen()) {
                        session.sendMessage(new TextMessage(jsonMessage));
                    }
                }
            } catch (Exception e) {
                System.err.println("发送广播消息失败: " + e.getMessage());
            }
        }
    }
    
    // 处理心跳
    private void handleHeartbeat(WebSocketSession session) {
        try {
            Map<String, Object> heartbeatResponse = new HashMap<>();
            heartbeatResponse.put("type", "HEARTBEAT_RESPONSE");
            sendJsonMessage(session, heartbeatResponse);
        } catch (Exception e) {
            System.err.println("发送心跳响应失败: " + e.getMessage());
        }
    }
    
    // 发送JSON消息
    private void sendJsonMessage(WebSocketSession session, Map<String, Object> message) throws Exception {
        String jsonMessage = objectMapper.writeValueAsString(message);
        session.sendMessage(new TextMessage(jsonMessage));
    }

    // 向多个特定用户发送消息
    public void sendToUsers(List<String> userIds, String message) {
        for (String userId : userIds) {
            sendToUser(userId, message);
        }
    }

    // 从会话中提取用户ID
    private String getUserIdFromSession(WebSocketSession session) {
        // 方式1: 从查询参数获取
        String query = session.getUri().getQuery();
        if (query != null && query.contains("userId=")) {
            try {
                // 正确解析URL参数，处理可能的多个参数情况
                String[] params = query.split("&");
                for (String param : params) {
                    if (param.startsWith("userId=")) {
                        return param.substring(7);
                    }
                }
            } catch (Exception e) {
                System.err.println("解析用户ID失败: " + e.getMessage());
            }
        }

        // 方式2: 从会话属性获取
        if (session.getAttributes().containsKey("userId")) {
            return (String) session.getAttributes().get("userId");
        }

        // 方式3: 生成匿名ID
        return "anonymous_" + session.getId().substring(0, 8);
    }

    // 获取当前连接数
    public static int getConnectionCount() {
        return sessions.size();
    }

    // 获取所有在线用户ID
    public static List<String> getOnlineUsers() {
        return new ArrayList<>(sessions.keySet());
    }
}
